#!/bin/bash

# Copyright 2019 Mobvoi Inc. All Rights Reserved.

. ./path.sh || exit 1;

# Use this to control how many gpu you use, It's 1-gpu training if you specify
# just 1gpu, otherwise it's is multiple gpu training based on DDP in pytorch
export CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7"
stage=5 # start from 0 if you need to start from data preparation
stop_stage=5

# You should change the following two parameters for multiple machine training,
# see https://pytorch.org/docs/stable/elastic/run.html
HOST_NODE_ADDR="localhost:0"
num_nodes=1


# data
data_url=www.openslr.org/resources/12
# use your own data path
datadir=/home/work_nfs2/xshi/corpus
# wav data dir
wave_data=/home/work_nfs7/bsmu/SALMONN
# Optional train_config
# 1. conf/train_transformer_large.yaml: Standard transformer
train_config=conf/train_salmonn_v7.yaml
# checkpoint=/home/local_data/bsmu/exp/salmonn_v1/1_27500.pt
checkpoint=
cmvn=true
do_delta=false

dir=/home/work_nfs7/xlgeng/bsmu_template/exp/salmonn_v8_lr5e_5
mkdir -p $dir
data_type=shard # raw
num_utts_per_shard=5000
# use average_checkpoint will get better result
average_checkpoint=false
decode_checkpoint=$dir/1_145000.pt
decode_checkpoint_name="1_145000"
# maybe you can try to adjust it if you can not get close results as README.md
average_num=10
# decode_modes="attention_rescoring ctc_greedy_search ctc_prefix_beam_search attention"
decode_modes="salmonn_decode"

. tools/parse_options.sh || exit 1;

# bpemode (unigram or bpe)
nbpe=5000
bpemode=unigram

set -e
set -u
set -o pipefail

train_set=train_clean_100
dev_set=dev

# test_sets="SPEECHIO_ASR_ZH00004"
test_sets=("aishell1" "aishell2" "SPEECHIO_ASR_ZH00000" "SPEECHIO_ASR_ZH00001" "SPEECHIO_ASR_ZH00002" "SPEECHIO_ASR_ZH00003" "SPEECHIO_ASR_ZH00004" "SPEECHIO_ASR_ZH00005" "test_meeting" "test_net")
shard_sets=("test_huawei_accent")
added_set_dir="/home/work_nfs4_ssd/xpyan/work_nfs5_ssd/HuaWei_SE_Android/wenet_meeting/data"
added_sets=("103" "assistant"  "xiaotiancai" )

dict=$wave_data/lang_char/${train_set}_${bpemode}${nbpe}_units.txt
bpemodel=$wave_data/lang_char/${train_set}_${bpemode}${nbpe}

if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
  echo 'decode by gengxuelong , start do decode to added set'
  # Test model, please specify the model you want to test by --checkpoint
  cmvn_opts=
  $cmvn && cmvn_opts="--cmvn data/${train_set}/global_cmvn"
  # TODO, Add model average here
  # mkdir -p $dir/test
  # if [ ${average_checkpoint} == true ]; then
  #   decode_checkpoint=$dir/avg_${average_num}.pt
  #   echo "do model average and final checkpoint is $decode_checkpoint"
  #   python wenet/bin/average_model.py \
  #     --dst_model $decode_checkpoint \
  #     --src_path $dir  \
  #     --num ${average_num} \
  #     --val_best
  # fi
  # Specify decoding_chunk_size if it's a unified dynamic chunk trained model
  # -1 for full chunk
  decoding_chunk_size=
  ctc_weight=0.5
  # Polling GPU id begin with index 0
  for test_set in "${added_sets[@]}"; do
  {
    echo "test this dataset: $test_set"
    test_set_name=$(echo "$test_set" | sed 's/\//_/g')
    test_dir=$dir/test_${decode_checkpoint_name}/${test_set_name}
    wer_path=$test_dir/wer
    if [ -e "$wer_path" ]; then
      echo 耿雪龙:"$wer_path 文件已存在，跳过对该数据集的推理"
      continue
    fi
    # if [ "$test_set" = "107" ]; then
    #   echo "耿雪龙:107数据集出现了NaN,不予推理"
    #   continue
    # fi
    mkdir -p $test_dir
    export CUDA_VISIBLE_DEVICES="7"
    python wenet/bin/recognize.py --gpu 7 \
      --mode $decode_modes \
      --config $dir/train.yaml \
      --data_type raw \
      --dict $dict \
      --bpe_model ${bpemodel}.model \
      --test_data $added_set_dir/$test_set/data.list \
      --checkpoint $decode_checkpoint \
      --beam_size 10 \
      --batch_size 1 \
      --penalty 0.0 \
      --result_file $test_dir/text_hyp \
      --ctc_weight $ctc_weight \
      # --test_data data/test/$test_set/data.list \

    # cut -f2- -d " " $test_dir/text_bpe > $test_dir/text_bpe_value_tmp
    # cut -f1 -d " " $test_dir/text_bpe > $test_dir/text_bpe_key_tmp
    # tools/spm_decode --model=${bpemodel}.model --input_format=piece \
    #   < $test_dir/text_bpe_value_tmp | sed -e "s/▁/ /g" > $test_dir/text_value_tmp
    # paste -d " " $test_dir/text_bpe_key_tmp $test_dir/text_value_tmp > $test_dir/text
    python tools/compute-wer.py --char=1 --v=1 \
      $added_set_dir/$test_set/text $test_dir/text_hyp > $test_dir/wer
    echo "$test_set has been decoded!"
  }
  done
  wait

fi
